from matplotlib import pyplot as plt

def smooth(read_path, file_name, x='timestep', y='reward', weight=0.75):
    # 平滑处理，类似tensorboard的smoothing函数
    fp = open(read_path+file_name, 'r')
    line = fp.readline()
    data = eval(line)
    scalar = data
    last = scalar[0]
    smoothed = []
    for point in scalar:
        smoothed_val = last * weight + (1 - weight) * point
        smoothed.append(smoothed_val)
        last = smoothed_val
    fp = open(read_path + 'smooth_' + file_name, 'w')
    print(smoothed, file=fp)
    fp.close()


def draw_pic(path, file_name, fig_path):
    fp = open(path + file_name)
    line = fp.readline()
    data = eval(line)
    plt.plot(data, color='lightgrey', label='data')

    fp = open(path + 'smooth_' + file_name)
    line = fp.readline()
    data = eval(line)
    # plt.plot(data, color=(255/255, 69/255, 0), label='smoothed')
    plt.plot(data, label='smoothed')

    plt.xlabel('counter')
    plt.ylabel('episode reward')
    # plt.ylabel('time (Engine clock)')
    # plt.ylabel('energy (Engine clock * W)')
    # plt.ylabel('loss')
    plt.legend(loc='upper right')
    
    plt.savefig(fig_path + file_name + '.png')
    plt.clf()

if __name__ == '__main__':
    ================= DQN ==========================
    smooth('./DQN/version1_result/data/', 'reward')
    draw_pic('./DQN/version1_result/data/', 'reward', './DQN/version1_result/figure/')

    smooth('./DQN/data/', 'time')
    draw_pic('./DQN/data/', 'time', './DQN/figure/')

    smooth('./DQN/data/', 'energy')
    draw_pic('./DQN/data/', 'energy', './DQN/figure/')

    smooth('./DQN/data/', 'loss')
    draw_pic('./DQN/data/', 'loss', './DQN/figure/')

    ================= PPO =============================
    smooth('./PPO/data/', 'reward')
    draw_pic('./PPO/data/', 'reward', './PPO/figure/')


    # ================ analysis ========================
    fp = open('./DQN/version1_result/data/smooth_reward')
    line = fp.readline()
    data = eval(line)
    plt.plot(data, label='original')

    fp = open('./DQN/version2_result/data/smooth_reward')
    line = fp.readline()
    data = eval(line)
    plt.plot(data, color=(255/255, 69/255, 0), label='normalized_state')
    plt.xlabel('counter')
    plt.ylabel('episode reward')
    plt.legend(loc='lower right')
    plt.savefig('./DQN/figure/compare_reward.png')
    plt.clf()